from datetime import datetime
import talib as ta
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import time
import redis
from math import *
import traceback
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import HistoryRequest, LogData
from vnpy.trader.rqdata import rqdata_client
from vnpy.trader.database import database_manager
from vnpy.trader.ui import create_qapp, QtCore

def handle_close(evt):
    print("Closed Figure!")
    exit(0)

FLAT_NUM= 30
SMA_NUM= 5

if __name__ == "__main__":
    
    symbol = "IF2201"
    exchange = Exchange.CFFEX
    start = datetime(2021, 12, 1)
    end = datetime(2022, 3, 17)
    interval = Interval.MINUTE
    bars = database_manager.load_bar_data(
        symbol, Exchange.CFFEX, interval=Interval.MINUTE, start=start, end=end
    )
    if not bars or len(bars) == 0:
        req = HistoryRequest(
            symbol=symbol,
            exchange=exchange,
            interval=Interval.MINUTE,
            start=start,
            end=end,
        )
        try:
            bars = rqdata_client.query_history(req)

            if bars:
                database_manager.save_bar_data(bars)
                print(f"{symbol}-{interval}历史数据下载完成")
            else:
                print(f"数据下载失败，无法获取{symbol}的历史数据")
        except Exception:
            msg = f"数据下载失败，触发异常：\n{traceback.format_exc()}"
            print(msg)
    n = FLAT_NUM
    history = bars[:n]
    left = len(bars) - n
    new_data = bars[n:]
    df = pd.DataFrame(
        [
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            }
            for s in history
        ]
    )

    fig = plt.figure()
    fig.canvas.mpl_connect("close_event", handle_close)

    last_tend = tend_value = 0
    offset_volume = 0
    pos = 0
    r = redis.Redis(host='localhost', port=6379, decode_responses=True)
    key = symbol
    nBalance = 100000000
    r.set(key ,  nBalance)  # 设置 name 对应的值
    last = df.iloc[-1:]
    print( f"\nBalance = %s" %(nBalance))
    s = last
    for i in range(len(new_data)):
        if not new_data or len(new_data) < 1:
            break
        s = new_data.pop(0)

        df = df.append(
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            },
            ignore_index=True,
        )

        plt.clf()  # 清空画布上的所有内容
        df["MeanSlow"] = ta.SMA(df["Close"], FLAT_NUM)
        df["MeanFast"] = ta.SMA(df["Close"], SMA_NUM)

        df["MeanVolume"] = ta.SMA(df["Volume"], FLAT_NUM)
        df["Close2"] = df.apply(lambda row: float(row["Close"]) - 4900, axis=1)

        df["weight"] = df.apply(
            lambda row: '%.2f' % ((row["Volume"] + 1) / (row["MeanVolume"] + 1)), axis=1
        )
        df["FlatS"] = df.apply(
            lambda row: (float(row["Close"]) - float(row["MeanSlow"]) *  float(row['weight'])) , axis=1
        )
        df["FlatF"] = df.apply(
            lambda row: (float(row["Close"]) - float(row["MeanFast"]) *  float(row['weight'])), axis=1
        )
        df["tends"] = ta.SMA( df["FlatS"] , FLAT_NUM)
        df["tendf"] = ta.SMA( df["FlatF"], SMA_NUM)
        # df['tends'] = df.apply(lambda row:  float(row['tends']) *  float(row['weight'])  ,axis =1)
        # df["tendf"] = df.apply(lambda row:  float(row['tendf']) *  float(row['weight'])  ,axis =1)
        df["one"] = 0
        plt.plot(df.index, df[["Close2", "tends","tendf",  "one"]])
        agNum = -SMA_NUM
        last = df.iloc [-1:]
        # last_tend = df.iloc[agNum:]
        tend_value = '%.2f' % float(last['tends'])
        last_tend = '%.2f' % float(last['tendf'])
        # last_tend = '%.2f' %  float(last_tend['tend'])
        # if ( float(last['Flat']) *(float(last['Flat']) * int(last['MeanVolume']) + (float(last['Close'])- float(last['Open'])*float(last["weight"])))) <0 :
        if float(last_tend )*float(tend_value) <0 :
            offset_volume = 0
            #rising buy
            if float(tend_value) >0 and pos <= 0:
                    offset_volume = 1 - pos
                    # print(f'buy  {offset_volume}')
            #desing shell

            elif float(tend_value) <0 and pos >= 0:
                    offset_volume = -1 - pos
                    # print(f'shell  {offset_volume}')
            else:
                # print('No condition')
                continue
            if offset_volume == 0 :
                continue
            money = -int(s.close_price *1000* offset_volume)
            nBalance = r.incrby(key , money)  # 设置 name 对应的值

            pos = pos + offset_volume
            last_tend = tend_value
            print(str(s.datetime.strftime('%m-%d%H:%M')), money, nBalance, "pos=%d orderv=%d weight=%s tend=%s Close2=%.1f Close=%s" %(pos,offset_volume, float(last['weight']), tend_value ,last['Close2']  ,float(last['Close'])))
        plt.pause(0.001)

    if pos != 0:
        money = int(s.close_price *3000 * pos)
        nBalance = r.incrby(key , money)  # 设置 name 对应的值
    print("Final%s Balance=%s pos = %d close = %s" %(str(s.datetime.strftime('%m-%d%H:%M')), nBalance, pos,  float(last['Close'])))

    # plt.plot(df.index, df[["Close2", "tend", "one"]])
    plt.show()
